#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据分类与统计报告生成器

功能：
1. 从评估结果中筛选有效数据
2. 按 evaluation.category 将数据分为三类
3. 生成标签分布统计报告

分类标准：
- auto_pass（自动通过）：票数≥80% 
- weak_consensus（弱一致）：默认N=5时，票数=3
- confusing（易混淆）：票数≤2 或未能形成有效投票
"""

import json
import os
import argparse
from typing import Dict, List, Tuple
from collections import Counter
import datetime

def load_jsonl_data(file_path: str) -> List[Dict]:
    """加载JSONL文件数据"""
    data = []
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                if line.strip():
                    try:
                        data.append(json.loads(line.strip()))
                    except json.JSONDecodeError as e:
                        print(f"第{line_num}行JSON解析错误: {e}")
                        continue
        print(f"成功加载 {len(data)} 条记录")
        return data
    except FileNotFoundError:
        print(f"文件不存在: {file_path}")
        return []
    except Exception as e:
        print(f"读取文件失败: {e}")
        return []

def filter_valid_data(data: List[Dict]) -> List[Dict]:
    """筛选有效数据（保持原有逻辑：只要有效且有label即算有效）"""
    valid_data = []
    
    for record in data:
        evaluation = record.get('evaluation', {})
        if (evaluation.get('is_valid', False) and 
            evaluation.get('label_id') is not None and
            evaluation.get('label_name') is not None):
            valid_data.append(record)
    
    invalid_count = len(data) - len(valid_data)
    print(f"有效数据: {len(valid_data)} 条")
    print(f"无效数据: {invalid_count} 条")
    if len(data) > 0:
        print(f"有效率: {len(valid_data)/len(data)*100:.1f}%")
    
    return valid_data

def split_by_category(valid_data: List[Dict]) -> Tuple[List[Dict], List[Dict], List[Dict]]:
    """按 evaluation.category 分为 auto_pass / weak_consensus / confusing"""
    categories = {'auto_pass': [], 'weak_consensus': [], 'confusing': []}
    for record in valid_data:
        category = record.get('evaluation', {}).get('category')
        if category in categories:
            categories[category].append(record)
    return categories['auto_pass'], categories['weak_consensus'], categories['confusing']

def analyze_label_distribution(valid_data: List[Dict]) -> Tuple[Dict[str, int], Dict[str, str]]:
    """分析标签分布"""
    label_counter = Counter()
    label_id_to_name: Dict[str, str] = {}

    for record in valid_data:
        evaluation = record.get('evaluation', {})
        label_id = evaluation.get('label_id')
        label_name = evaluation.get('label_name')
        if label_id:
            key = str(label_id)
            label_counter[key] += 1
            # 以首次出现的名称为准
            if key not in label_id_to_name and label_name:
                label_id_to_name[key] = str(label_name)

    return dict(label_counter), label_id_to_name

def save_valid_data(valid_data: List[Dict], output_file: str):
    """保存数据到新的JSONL文件"""
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            for record in valid_data:
                f.write(json.dumps(record, ensure_ascii=False) + '\n')
        print(f"数据已保存到: {output_file}")
    except Exception as e:
        print(f"保存文件失败: {e}")

def generate_comprehensive_label_report(overall_distribution: Dict[str, int], auto_pass_distribution: Dict[str, int], 
                                      output_file: str, label_id_to_name: Dict[str, str]):
    """生成综合标签分布报告（总体和auto_pass类别分别统计）"""
    overall_total = sum(overall_distribution.values())
    auto_pass_total = sum(auto_pass_distribution.values())

    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write("标签分布统计报告\n")
            f.write(f"生成时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"{'='*60}\n\n")
            
            # 第一部分：总体标签分布
            f.write("## 1. 总体标签分布（所有有效数据）\n")
            f.write(f"总样本数: {overall_total}\n\n")
            
            if overall_distribution:
                for label_id in sorted(overall_distribution.keys(), key=lambda x: int(x)):
                    count = overall_distribution[label_id]
                    percentage = count / overall_total * 100 if overall_total > 0 else 0
                    label_name = label_id_to_name.get(label_id, "未知标签")
                    f.write(f"标签 {label_id}: {label_name}\n")
                    f.write(f"  数量: {count} ({percentage:.1f}%)\n\n")
            else:
                f.write("  无总体标签数据\n\n")
            
            f.write(f"{'='*60}\n\n")
            
            # 第二部分：auto_pass标签分布
            f.write("## 2. auto_pass类别标签分布（高质量数据）\n")
            f.write(f"总样本数: {auto_pass_total}\n\n")
            
            if auto_pass_distribution:
                for label_id in sorted(auto_pass_distribution.keys(), key=lambda x: int(x)):
                    count = auto_pass_distribution[label_id]
                    percentage = count / auto_pass_total * 100 if auto_pass_total > 0 else 0
                    label_name = label_id_to_name.get(label_id, "未知标签")
                    f.write(f"标签 {label_id}: {label_name}\n")
                    f.write(f"  数量: {count} ({percentage:.1f}%)\n\n")
            else:
                f.write("  无auto_pass标签数据\n\n")

        print(f"标签分布报告已保存到: {output_file}")
    except Exception as e:
        print(f"生成报告失败: {e}")

def main():
    parser = argparse.ArgumentParser(description='数据分类与统计报告生成器')
    parser.add_argument('input_file', help='输入的JSONL文件路径')
    parser.add_argument('--label-data-dir', default='../data/label_data', help='三类数据输出目录')
    parser.add_argument('--log-dir', default='../logs/label_log', help='统计报告输出目录')
    parser.add_argument('--save-valid', action='store_true', help='是否保存有效数据到新文件（自动通过）')
    parser.add_argument('--save-weak', action='store_true', help='是否保存弱一致数据到新文件')
    parser.add_argument('--save-confusing', action='store_true', help='是否保存易混淆数据到新文件')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.label_data_dir, exist_ok=True)
    os.makedirs(args.log_dir, exist_ok=True)
    
    # 加载数据
    print(f"正在加载数据: {args.input_file}")
    data = load_jsonl_data(args.input_file)
    
    if not data:
        print("没有加载到任何数据，程序退出")
        return
    
    # 筛选有效数据
    print("\n正在筛选有效数据...")
    valid_data = filter_valid_data(data)
    
    if not valid_data:
        print("没有找到有效数据，程序退出")
        return
    
    # 生成时间戳
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    
    # 按类别分割
    auto_pass, weak_consensus, confusing = split_by_category(valid_data)
    
    # 分析标签分布（总体和auto_pass分别统计）
    print("\n正在分析标签分布...")
    print("  - 统计总体有效数据的标签分布")
    overall_label_distribution, overall_label_id_to_name = analyze_label_distribution(valid_data)
    print("  - 统计auto_pass类别的标签分布")
    auto_pass_label_distribution, auto_pass_label_id_to_name = analyze_label_distribution(auto_pass)
    
    # 合并标签名称映射（以总体为准，auto_pass补充）
    label_id_to_name = overall_label_id_to_name.copy()
    label_id_to_name.update(auto_pass_label_id_to_name)

    # 保存数据（可选）
    if args.save_valid:
        valid_data_file = os.path.join(args.label_data_dir, f'auto_pass_{timestamp}.jsonl')
        save_valid_data(auto_pass, valid_data_file)
    if args.save_weak:
        weak_file = os.path.join(args.label_data_dir, f'weak_consensus_{timestamp}.jsonl')
        save_valid_data(weak_consensus, weak_file)
    if args.save_confusing:
        confusing_file = os.path.join(args.label_data_dir, f'confusing_{timestamp}.jsonl')
        save_valid_data(confusing, confusing_file)
    
    # 生成标签分布报告
    report_file = os.path.join(args.log_dir, f'label_distribution_{timestamp}.txt')
    generate_comprehensive_label_report(overall_label_distribution, auto_pass_label_distribution, report_file, label_id_to_name)
    
    # 打印简要统计
    print(f"\n简要统计:")
    print(f"总样本数: {len(data)}")
    print(f"有效样本数: {len(valid_data)} (auto_pass={len(auto_pass)}, weak={len(weak_consensus)}, confusing={len(confusing)})")
    print(f"总体标签种类数: {len(overall_label_distribution)}")
    print(f"auto_pass类别标签种类数: {len(auto_pass_label_distribution)}")
    print(f"三类数据保存在: {args.label_data_dir}")
    print(f"统计报告保存在: {args.log_dir}")

if __name__ == "__main__":
    main() 